package com.milvus.service.impl;

import com.milvus.service.SearcherService;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.SearchResp;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@Service
public class SearcherServiceImpl extends MilvusServiceImpl implements SearcherService {

    public SearcherServiceImpl() throws Exception {
        super();
    }
    public List<List<SearchResp.SearchResult>> search(MultipartFile multipartFile) throws IOException {
        File file = multipartFileToFile(multipartFile);
        try {

            INDArray queryFeatures = extractor.extractFeatures(file);


            // 在Milvus中进行相似性搜索
            float[] floatArray = queryFeatures.toFloatVector();
            List<BaseVector> floatList = new ArrayList<>();

            FloatVec floatVec = new FloatVec(floatArray);
            floatList.add(floatVec);

            SearchReq searchReq = SearchReq.builder()
                    .collectionName(collectionName)
                    .data(floatList)
                    .outputFields(Arrays.asList("name"))
                    .topK(10)
                    .build();

            SearchResp searchResp = milvusClient.search(searchReq);

            List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
            return searchResults;
        }catch (Exception e){

        }finally {
            file.delete();
        }
        return null;
    }
    public File multipartFileToFile(MultipartFile multipartFile){
        File file = new File(uploadPath+multipartFile.getName());
        try {
            multipartFile.transferTo(file);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return file;
    }
}
